'''
Created on Nov 17, 2010

@author: oabalbin
'''

'''
time /exds/sw/bioinfo/alignment/samtools/samtools-0.1.8/samtools 
mpileup -gf /exds/projects/alignment_indexes/bwa/hg19_illumina/hg19_illumina.fa ../aM18/s_3_12_sequence.aln.rmdup.sorted.bam.bam ../aM18/s_4_12_sequence.aln.rmdup.sorted.bam.bam ../aM27/s_5_12_sequence.aln.rmdup.sorted.bam.bam ../aM27/s_6_12_sequence.aln.rmdup.sorted.bam.bam ../aM44/s_7_12_sequence.aln.rmdup.sorted.bam.bam ../aM44/s_8_12_sequence.aln.rmdup.sorted.bam.bam > all_tumors_plus_normals.aln.rmdup.bcf
'''
import os
import glob
import multiprocessing
import subprocess
from snps.configfile import config_run, config_experiment
from collections import deque

'''
def check_bfrm_output(self, bfrmfolder,newFFile=[]):
     """
     Check if files  mA.txt, mF.txt, mPostPib.txt, exist
     """
     for f in bfrmof:
         fname = bfrmfolder + f
         print fname
         if os.path.isfile(fname):
             files.append(fname)

     return files
'''

def read_files_folder(folderpath,ext):
    ''' '''
    # Read files in folder
    myfiles=[]
    for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
        myfiles.append(infile)
        #print "current file is: " + infile
    return myfiles


def samtools_pileup(thargs, bamfiles, outfile):
    '''
    Returns the shell to execute a pile up using sam tools 
    param is a dictionary with the input param
    '''
    args=['samtools','mpileup',thargs['mpileup_args'],thargs['ref.fa']] + bamfiles
    
    f = open(outfile, "w")
    retcode = subprocess.call(args, stdout=f)
    f.close()
    
    # You can use the returncode for checking the process status
    
    return outfile

def bcftools_call_snps(bcf_file_name):
    '''
    Call snps candidates using bcftools
    This tool converts BCF to VCF, indexes BCF for random access, concatenates (not merges) BCFs, 
    estimate site allele frequencies and calls SNP candidates.
    '''
    outfile_name=bcf_file_name.replace('.bcf','.vcf')
    args = ['bcftools', 'view', '-vcg', bcf_file_name]
    f = open(outfile_name, "w")
    retcode = subprocess.call(args, stdout=f)
    f.close()
    
    return outfile_name



##### vcf files analysis using vcftools
'''
/exds/sw/bioinfo/alignment/vcftools/cpp/vcftools --vcf s_3_12_sequence.aln.rmdup.candidates.vcf 
--minQ 20 --out /exds/users/oabalbin/projects/snps/exomes/aM18/s_3_12_sequence.aln.rmdup.candidatesR 
--recode --keep-INFO DP --keep-INFO AF1 --keep-INFO MQ --keep-INFO CI95 --keep-INFO INDEL
You can also filter by mean depth of coverage an other options, look at the vcf tools.
'''


def apply_snps_filters(vcf_file_name, filters_dict):
    '''
    Apply filters to the vcf file in order to generate the final
    call of snps. 
    '''
    filters=[]
    for thf, value in filters_dict.iteritems():
        filters.append(thf)
        filters.append(value)
        
    outfile_name = vcf_file_name.replace('.vcf','')
    # Note: remove the path when the vcftools added to the general path
    args=['/exds/sw/bioinfo/alignment/vcftools/cpp/vcftools','--vcf',vcf_file_name] + filters + \
    ['--out',outfile_name,'--recode', '--keep-INFO', 'DP', 
     '--keep-INFO', 'AF1', '--keep-INFO', 'MQ','--keep-INFO', 
     'CI95', '--keep-INFO', 'INDEL']
    
    retcode = subprocess.call(args)
    return outfile_name+'.recode.vcf'

    
####################
def vfc_validation(inputfile_name):
    '''
    It validates the vcf input file.
    It is good to check if the vcf files are ok after calling the snps with sam-tools
    It is not necessary when running other vcf-tools because they do also the check 
    '''
    args=['vcf-validator',inputfile_name]
    retcode = subprocess.call(args)
    return retcode



def snps_stats(inputfile_name):
    '''
    '''
    outfile_name=inputfile_name.replace('.vcf','.vcf.stats')
    args = ['vcf-stats', inputfile_name]
    f = open(outfile_name, "w")
    retcode = subprocess.call(args, stdout=f)
    f.close()
    
    return retcode

def vcf2tab(inputfile_name):
    '''
    It writes the vcf file into a tab delimited format with the actual variants, not all the ALT
    If the files are gzip=> zcat, otherwise cat
    '''
    outfile_name=inputfile_name.replace('.vcf','.tab')
    args = ['cat', inputfile_name,'|','inputfile_name']
    f = open(outfile_name, "w")
    retcode = subprocess.call(args, stdout=f)
    f.close()
    
    return retcode


def shared_snps_btw_samples(inpufile_list,pathdir):
    '''
    Creates intersections and complements of two or more VCF files. 
    Given multiple VCF files, it can output the list of positions which are shared by at least N files, at most N files, exactly N files, etc. 
    It receives a list of input files
    '''
    
    outfile_name=pathdir+'.vcf.isec.gz'
    
    args = ['vcf-isec', '-n', '=', len(inpufile_list)]+ inpufile_list + ['|','bgzip','-c']
    f = open(outfile_name, "w")
    retcode = subprocess.call(args, stdout=f)
    f.close()
    
    return retcode

       

def main(configrun, configexp,snpcall_method):
    run_param = config_run(configrun)
    experiment_run = config_experiment(configexp)
    #ext='.rmdup.sorted.bam.bam'
    ext='.sorted.bam'
    outfile='/exds/users/oabalbin/projects/snps/exomes/test/test_samfile.bcf'
    experiment_run = experiment_run['SampleDir']
    spbamfiles=[]
    
    num_pileups=len(experiment_run)
    print 'Creating pool with %d processes\n' % num_pileups 
    pool = multiprocessing.Pool(num_pileups)
    print 'pool = %s' % pool
    
    ####
    for sp,fdpath in experiment_run.iteritems():
        spbamfiles.append(read_files_folder(fdpath,ext))
        
        spbamfiles=sum(spbamfiles,[])
    
    # Do the mpileup    
    bcf_file_name = samtools_pileup(run_param[snpcall_method], spbamfiles, outfile)
    # Create the vcf file of candidate snps
    vcf_file_name = bcftools_call_snps(bcf_file_name)
    
    # Filter candidate snps according to qualtity and coverage depth
    # Filter_dict should be in the config file or hardcode here
    filters_dict={'--minQ': '20'}
    filtered_vcf_file = apply_snps_filters(vcf_file_name, filters_dict)
    
    
    return filtered_vcf_file # Return the location of the filtered snps file.
    
    
    
    
    
    
    
    
    